3b7803
@@ -1,11 +1,11 @@
 /*
- * Copyright (c) 2011 by the original author(s).
+ * Copyright 2011-2012 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
  *
- *     http://www.apache.org/licenses/LICENSE-2.0
+ *      http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
@@ -19,7 +19,6 @@
import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 
-import org.bson.types.BasicBSONList;
 import org.bson.types.ObjectId;
 import org.springframework.core.convert.ConversionException;
 import org.springframework.core.convert.ConversionService;
@@ -35,11 +34,12 @@
import org.springframework.util.Assert;
 import com.mongodb.BasicDBList;
 import com.mongodb.BasicDBObject;
 import com.mongodb.DBObject;
+import com.mongodb.DBRef;
 
 /**
  * A helper class to encapsulate any modifications of a Query object before it gets submitted to the database.
  * 
- * @author Jon Brisbin <jbrisbin@vmware.com>
+ * @author Jon Brisbin
  * @author Oliver Gierke
  */
 public class QueryMapper {
@@ -75,8 +75,8 @@
public class QueryMapper {
 	 */
 	public DBObject getMappedObject(DBObject query, MongoPersistentEntity<?> entity) {
 
-		if (isKeyWord(query)) {
-			return getMappedKeyword(query, entity);
+		if (Keyword.isKeyword(query)) {
+			return getMappedKeyword(new Keyword(query), entity);
 		}
 
 		DBObject result = new BasicDBObject();
@@ -100,25 +100,38 @@
public class QueryMapper {
 	 * @param entity
 	 * @return
 	 */
-	private DBObject getMappedKeyword(DBObject query, MongoPersistentEntity<?> entity) {
-
-		String newKey = query.keySet().iterator().next();
-		Object value = query.get(newKey);
+	private DBObject getMappedKeyword(Keyword query, MongoPersistentEntity<?> entity) {
 
 		// $or/$nor
-		if (newKey.matches(N_OR_PATTERN)) {
+		if (query.key.matches(N_OR_PATTERN)) {
 
-			Iterable<?> conditions = (Iterable<?>) value;
+			Iterable<?> conditions = (Iterable<?>) query.value;
 			BasicDBList newConditions = new BasicDBList();
 
 			for (Object condition : conditions) {
 				newConditions.add(getMappedObject((DBObject) condition, entity));
 			}
 
-			return new BasicDBObject(newKey, newConditions);
+			return new BasicDBObject(query.key, newConditions);
 		}
 
-		return new BasicDBObject(newKey, convertSimpleOrDBObject(value, entity));
+		return new BasicDBObject(query.key, convertSimpleOrDBObject(query.value, entity));
+	}
+
+	/**
+	 * Returns the mapped keyword considered defining a criteria for the given property.
+	 * 
+	 * @param keyword
+	 * @param property
+	 * @return
+	 */
+	public DBObject getMappedKeyword(Keyword keyword, MongoPersistentProperty property) {
+
+		if (property.isAssociation()) {
+			convertAssociation(keyword.value, property);
+		}
+
+		return new BasicDBObject(keyword.key, getMappedValue(keyword.value, property, keyword.key));
 	}
 
 	/**
@@ -161,7 +174,7 @@
public class QueryMapper {
 		}
 
 		if (property.isAssociation()) {
-			return isKeyWord(source) ? getMappedValue(getKeywordValue(source), property, newKey) : convertAssociation(source,
+			return Keyword.isKeyword(source) ? getMappedKeyword(new Keyword(source), property) : convertAssociation(source,
 					property);
 		}
 
@@ -243,14 +256,14 @@
public class QueryMapper {
 		}
 
 		if (source instanceof Iterable) {
-			BasicBSONList result = new BasicBSONList();
+			BasicDBList result = new BasicDBList();
 			for (Object element : (Iterable<?>) source) {
-				result.add(converter.toDBRef(element, property));
+				result.add(element instanceof DBRef ? element : converter.toDBRef(element, property));
 			}
 			return result;
 		}
 
-		return converter.toDBRef(source, property);
+		return source instanceof DBRef ? source : converter.toDBRef(source, property);
 	}
 
 	/**
@@ -276,47 +289,59 @@
public class QueryMapper {
 	}
 
 	/**
-	 * Returns whether the given value is representing a query keyword.
+	 * Converts the given raw id value into either {@link ObjectId} or {@link String}.
 	 * 
-	 * @param value
+	 * @param id
 	 * @return
 	 */
-	private static boolean isKeyWord(Object value) {
+	public Object convertId(Object id) {
 
-		if (!(value instanceof DBObject) || value instanceof BasicDBList) {
-			return false;
+		try {
+			return conversionService.convert(id, ObjectId.class);
+		} catch (ConversionException e) {
+			// Ignore
 		}
 
-		DBObject dbObject = (DBObject) value;
-		return dbObject.keySet().size() == 1 && dbObject.keySet().iterator().next().startsWith("$");
+		return converter.convertToMongoType(id);
 	}
 
 	/**
-	 * Returns the value of the given source assuming it's a query keyword.
+	 * Value object to capture a query keyword representation.
 	 * 
-	 * @param source
-	 * @return
+	 * @author Oliver Gierke
 	 */
-	private static Object getKeywordValue(Object source) {
+	private static class Keyword {
 
-		DBObject dbObject = (DBObject) source;
-		return dbObject.get(dbObject.keySet().iterator().next());
-	}
+		String key;
+		Object value;
 
-	/**
-	 * Converts the given raw id value into either {@link ObjectId} or {@link String}.
-	 * 
-	 * @param id
-	 * @return
-	 */
-	public Object convertId(Object id) {
+		Keyword(Object source) {
 
-		try {
-			return conversionService.convert(id, ObjectId.class);
-		} catch (ConversionException e) {
-			// Ignore
+			Assert.isInstanceOf(DBObject.class, source);
+
+			DBObject value = (DBObject) source;
+
+			Assert.isTrue(value.keySet().size() == 1, "Keyword must have a single key only!");
+
+			this.key = value.keySet().iterator().next();
+			this.value = value.get(key);
 		}
 
-		return converter.convertToMongoType(id);
+		/**
+		 * Returns whether the given value actually represents a keyword. If this returns {@literal true} it's safe to call
+		 * the constructor.
+		 * 
+		 * @param value
+		 * @return
+		 */
+		static boolean isKeyword(Object value) {
+
+			if (!(value instanceof DBObject)) {
+				return false;
+			}
+
+			DBObject dbObject = (DBObject) value;
+			return dbObject.keySet().size() == 1 && dbObject.keySet().iterator().next().startsWith("$");
+		}
 	}
 }
